# Munsell system of colors from Zouxy : http://blog.csdn.net/zouxy09
import numpy as np
import cv2
import math


UNKNOWN_FLOW_THRESH = 1e9

def make_color_wheel():
	lst_color_wheel = []

	RY = 15
	YG = 6
	GC = 4
	CB = 11
	BM = 13
	MR = 6

	for i in range(RY):
		lst_color_wheel.append(np.array([255, 255 * i / RY, 0], dtype = np.uint8))
	for i in range(YG):
		lst_color_wheel.append(np.array([255 - 255 * i / YG, 255, 0], dtype = np.uint8))
	for i in range(GC):
		lst_color_wheel.append(np.array([0, 255, 255 * i / GC], dtype = np.uint8))
	for i in range(CB):
		lst_color_wheel.append(np.array([0, 255 - 255 * i / CB, 255], dtype = np.uint8))
	for i in range(BM):
		lst_color_wheel.append(np.array([255 * i / BM, 0, 255], dtype = np.uint8))
	for i in range(MR):
		lst_color_wheel.append(np.array([255, 0, 255 - 255 * i / MR], dtype = np.uint8))

	return lst_color_wheel


def motion_to_color(flow):
	flow_color = np.zeros((flow.shape[0],flow.shape[1],3), dtype = np.uint8)
	lst_color_wheel = make_color_wheel()
	maxrad = -1.0

	for i in range(flow.shape[0]):
		for j in range(flow.shape[1]):
			fx = flow[i,j,0]
			fy = flow[i,j,1]
			if math.fabs(fx) > UNKNOWN_FLOW_THRESH or math.fabs(fy) > UNKNOWN_FLOW_THRESH:
				continue
			rad = math.sqrt(fx * fx + fy * fy)
			if maxrad <= rad:
				maxrad = rad
	if maxrad <= 0:
		return False, None

	for i in range(flow.shape[0]):
		for j in range(flow.shape[1]):
			n_p = 240 * i + 3 * j
			fx = flow[i,j,0] / maxrad
			fy = flow[i,j,1] / maxrad
			if math.fabs(fx) > UNKNOWN_FLOW_THRESH or math.fabs(fy) > UNKNOWN_FLOW_THRESH:
				continue

			rad = math.sqrt(fx * fx + fy * fy)
			angle = math.atan2(-fy, -fx) / math.pi
			fk = (angle + 1.0) / 2.0 * (len(lst_color_wheel) - 1)
			k0 = int(fk)
			k1 = int((k0 + 1) % len(lst_color_wheel))
			f = fk - k0

			for b in range(3):
				col0 = (lst_color_wheel[k0])[b] / 255.0
				col1 = (lst_color_wheel[k1])[b] / 255.0
				col = (1 - f) * col0 + f * col1
				if rad <= 1:
					col = 1 - rad * (1 - col)
				else:
					col = col * 0.75
				flow_color[i,j,2-b] = int(255.0 * col)

	return flow_color
